# 股票交易数据整理
from builtins import float

from sqlalchemy import create_engine, Column, Integer, String, Date, DateTime, Float
from sqlalchemy.orm import sessionmaker, declarative_base
import stock_data_get
from datetime import date

# 创建数据库引擎
engine = create_engine('mysql+pymysql://stock:stock@localhost:3306/stock?charset=utf8')

# 创建基类模型
Base = declarative_base()

# 创建交易记录类
class DealRecord(Base):
    __tablename__ = 'deal_record'
    id = Column(Integer, primary_key=True, autoincrement=True)
    code = Column(String(20), nullable=False)
    name = Column(String(50), nullable=False)
    deal_time = Column(DateTime, nullable=False)
    deal_price = Column(Float, nullable=False)
    deal_num = Column(Integer, nullable=False)
    deal_money = Column(Float, nullable=False)
    deal_type = Column(String(2), nullable=False)

# 创建交易统计类
class DealStatistic(Base):
    __tablename__ = 'deal_statistic'
    id = Column(Integer, primary_key=True, autoincrement=True)
    code = Column(String(20), nullable=False)
    name = Column(String(50), nullable=False)
    deal_date = Column(Date, nullable=False)
    start_price = Column(Float, nullable=False)
    end_price = Column(Float, nullable=False)
    high_price = Column(Float, nullable=False)
    low_price = Column(Float, nullable=False)
    deal_num = Column(Integer, nullable=False)
    deal_money = Column(Float, nullable=False)
    buy_num = Column(Integer, nullable=False)
    buy_money = Column(Float, nullable=False)
    sale_num = Column(Integer, nullable=False)
    sale_money = Column(Float, nullable=False)
    mid_num = Column(Integer, nullable=False)
    mid_money = Column(Float, nullable=False)

def get_and_deal_today_data(code, name, addr):
    # 清除今天的数据，避免重复导入
    deleteDealRecordAndStatistic(code, date.today())

    # 创建会话
    session = sessionmaker(bind=engine)()\

    # 处理交易记录数据
    data_list = []
    date_str = date.today().strftime("%Y-%m-%d")
    stock_data_get.get_all_data(data_list, addr)
    new_deal_record_list = []

    for one_data in data_list:
        one_data_list = one_data.split("/")
        new_deal_record = DealRecord(
            code=code,
            name=name,
            deal_time= date_str + " " + one_data_list[1],
            deal_price=float(one_data_list[2]),
            deal_num=int(one_data_list[4]),
            deal_money=float(one_data_list[5]),
            deal_type=one_data_list[6]
        )
        new_deal_record_list.append(new_deal_record)
        session.add(new_deal_record)

    # 处理交易统计数据
    start_price = new_deal_record_list[0].deal_price
    end_price = new_deal_record_list[-1].deal_price
    deal_price_list = [obj.deal_price for obj in new_deal_record_list]
    high_price = max(deal_price_list)
    low_price = min(deal_price_list)
    buy_num = 0
    buy_money = 0.0
    sale_num = 0
    sale_money = 0.0
    mid_num = 0
    mid_money = 0.0
    for record in new_deal_record_list:
        if record.deal_type == "B":
            buy_num += record.deal_num
            buy_money += record.deal_money
        else:
            if record.deal_type == "S":
                sale_num += record.deal_num
                sale_money += record.deal_money
            else:
                mid_num += record.deal_num
                mid_money += record.deal_money

    new_deal_statistic = DealStatistic(
        code=code,
        name=name,
        deal_date=date.today(),
        start_price=start_price,
        end_price=end_price,
        high_price=high_price,
        low_price=low_price,
        deal_num = buy_num + sale_num,
        deal_money = buy_money + sale_money,
        buy_num = buy_num,
        buy_money = buy_money,
        sale_num = sale_num,
        sale_money = sale_money,
        mid_num = mid_num,
        mid_money = mid_money
        )
    session.add(new_deal_statistic)

    # 插入数据
    session.commit()
    session.close()

def getDealRecord(code, date = None):
    session = sessionmaker(bind=engine)()
    if date is None:
        deal_record_list = session.query(DealRecord).filter(DealRecord.code == code).all()
    else:
        deal_record_list = session.query(DealRecord).filter(DealRecord.code == code, DealRecord.deal_time.like(date.strftime('%Y-%m-%d') + "%")).all()
    return deal_record_list

def getDealStatistic(code, date = None):
    session = sessionmaker(bind=engine)()
    if date is None:
        deal_statistic_list = session.query(DealStatistic).filter(DealStatistic.code == code).all()
    else:
        deal_statistic_list = session.query(DealStatistic).filter(DealStatistic.code == code, DealStatistic.deal_date == date).all()
    return deal_statistic_list

def deleteDealRecord(code, date = None):
    session = sessionmaker(bind=engine)()
    if date is None:
        session.query(DealRecord).filter(DealRecord.code == code).delete()
    else:
        session.query(DealRecord).filter(DealRecord.code == code, DealRecord.deal_time.like(date.strftime('%Y-%m-%d') + "%")).delete()
        # session.query(DealRecord).filter(DealRecord.code == code, DealRecord.deal_time.like(date + "%")).delete()
    session.commit()

def deleteDealStatistic(code, date = None):
    session = sessionmaker(bind=engine)()
    if date is None:
        session.query(DealStatistic).filter(DealStatistic.code == code).delete()
    else:
        session.query(DealStatistic).filter(DealStatistic.code == code, DealStatistic.deal_date == date).delete()
    session.commit()

def deleteDealRecordAndStatistic(code, date = None):
    deleteDealRecord(code, date)
    deleteDealStatistic(code, date)

# 采集目标数据
if __name__ == '__main__':
    code_name_dict_list = [{"code":"600825","name":"新华传媒","addr":"sh600825"},
                           {"code":"000100","name":"TCL科技","addr":"sz000100"},
                           {"code":"603501","name":"韦尔股份","addr":"sh603501"},
                           {"code":"601360","name":"三六零","addr":"sh601360"},
                           {"code":"002230","name":"科大讯飞","addr":"sz002230"}]
    for codeNameDict in code_name_dict_list:
        code = codeNameDict["code"]
        name = codeNameDict["name"]
        addr = codeNameDict["addr"]
        # 采集今天的数据
        get_and_deal_today_data(code, name, addr)
